import os
import sys
import importlib.util
import inspect
import json
import re
from typing import Dict, List, Any, Callable
from app.utils.logger import logger

class VulnerabilityScriptManager:
    """
    漏洞脚本管理器
    负责自动加载、注册和管理漏洞扫描脚本
    """
    
    def __init__(self, script_dir=None):
        # 默认脚本目录
        if script_dir is None:
            self.script_dir = os.path.join(os.path.dirname(os.path.dirname(os.path.dirname(__file__))), 'scripts')
        else:
            self.script_dir = script_dir
        
        # 脚本存储字典
        self.scripts: Dict[str, Dict[str, Any]] = {}
        
        # 脚本元数据存储
        self.script_metadata_file = os.path.join(self.script_dir, 'script_metadata.json')
        
        # 确保脚本目录存在
        os.makedirs(self.script_dir, exist_ok=True)
        
        # 初始化加载脚本
        self.load_scripts()
    
    def load_scripts(self):
        """
        加载所有漏洞扫描脚本
        """
        logger.info(f"开始加载漏洞扫描脚本，目录: {self.script_dir}")
        
        # 清除现有脚本
        self.scripts.clear()
        
        try:
            # 遍历脚本目录
            for filename in os.listdir(self.script_dir):
                # 只处理Python文件，跳过__init__.py和以_开头的文件
                if filename.endswith('.py') and not filename.startswith('_') and filename != '__init__.py':
                    script_path = os.path.join(self.script_dir, filename)
                    script_name = filename[:-3]  # 去掉.py扩展名
                    
                    try:
                        # 加载脚本
                        spec = importlib.util.spec_from_file_location(script_name, script_path)
                        if spec and spec.loader:
                            module = importlib.util.module_from_spec(spec)
                            sys.modules[script_name] = module
                            spec.loader.exec_module(module)
                            
                            # 分析脚本内容
                            script_info = self._analyze_script(module, script_path)
                            if script_info:
                                self.scripts[script_name] = script_info
                                logger.info(f"成功加载脚本: {script_name}")
                    except Exception as e:
                        logger.error(f"加载脚本 {filename} 失败: {str(e)}")
            
            # 保存元数据
            self._save_metadata()
            logger.info(f"脚本加载完成，共加载 {len(self.scripts)} 个脚本")
            
        except Exception as e:
            logger.error(f"加载脚本时发生错误: {str(e)}")
    
    def _analyze_script(self, module, script_path):
        """
        分析脚本内容，提取脚本信息
        """
        try:
            # 检查脚本是否包含必要的属性和函数
            if not hasattr(module, 'scan') or not callable(module.scan):
                logger.warning(f"脚本 {os.path.basename(script_path)} 缺少scan函数")
                return None
            
            # 提取脚本信息
            script_info = {
                'name': getattr(module, 'NAME', os.path.basename(script_path)[:-3]),
                'description': getattr(module, 'DESCRIPTION', '无描述'),
                'author': getattr(module, 'AUTHOR', '未知'),
                'version': getattr(module, 'VERSION', '1.0'),
                'vulnerability_type': getattr(module, 'VULNERABILITY_TYPE', 'unknown'),
                'severity': getattr(module, 'SEVERITY', 'medium'),
                'scan_function': module.scan,
                'path': script_path,
                'enabled': getattr(module, 'ENABLED', True),
                'required_params': getattr(module, 'REQUIRED_PARAMS', []),
                'optional_params': getattr(module, 'OPTIONAL_PARAMS', {})
            }
            
            # 获取scan函数的参数信息
            sig = inspect.signature(module.scan)
            script_info['function_params'] = list(sig.parameters.keys())
            
            return script_info
        except Exception as e:
            logger.error(f"分析脚本 {script_path} 失败: {str(e)}")
            return None
    
    def get_scripts(self, filter_enabled=True):
        """
        获取所有脚本列表
        filter_enabled: 是否只返回启用的脚本
        """
        if filter_enabled:
            return [script for script in self.scripts.values() if script.get('enabled', True)]
        return list(self.scripts.values())
    
    def get_script_by_name(self, script_name):
        """
        根据名称获取脚本
        """
        return self.scripts.get(script_name)
    
    def get_scripts_by_type(self, vuln_type):
        """
        根据漏洞类型获取脚本
        """
        return [script for script in self.scripts.values() 
                if script.get('vulnerability_type') == vuln_type and script.get('enabled', True)]
    
    def execute_script(self, script_name, target, **kwargs):
        """
        执行指定脚本
        """
        script = self.get_script_by_name(script_name)
        if not script:
            logger.error(f"脚本 {script_name} 不存在")
            return None
        
        if not script.get('enabled', True):
            logger.error(f"脚本 {script_name} 已禁用")
            return None
        
        try:
            logger.info(f"执行脚本: {script_name}, 目标: {target}")
            result = script['scan_function'](target, **kwargs)
            
            # 确保结果格式标准化
            if result and isinstance(result, dict):
                # 添加脚本元信息
                result['script_name'] = script_name
                result['vulnerability_type'] = script['vulnerability_type']
                result['severity'] = script.get('severity', 'medium')
                
            return result
        except Exception as e:
            logger.error(f"执行脚本 {script_name} 失败: {str(e)}")
            return None
    
    def create_script(self, script_name, content):
        """
        创建新脚本
        """
        if script_name in self.scripts:
            logger.error(f"脚本 {script_name} 已存在")
            return False
        
        try:
            script_path = os.path.join(self.script_dir, f"{script_name}.py")
            with open(script_path, 'w', encoding='utf-8') as f:
                f.write(content)
            
            # 重新加载脚本
            self.load_scripts()
            logger.info(f"脚本 {script_name} 创建成功")
            return True
        except Exception as e:
            logger.error(f"创建脚本 {script_name} 失败: {str(e)}")
            return False
    
    def update_script(self, script_name, content):
        """
        更新脚本
        """
        if script_name not in self.scripts:
            logger.error(f"脚本 {script_name} 不存在")
            return False
        
        try:
            script_path = self.scripts[script_name]['path']
            with open(script_path, 'w', encoding='utf-8') as f:
                f.write(content)
            
            # 重新加载脚本
            self.load_scripts()
            logger.info(f"脚本 {script_name} 更新成功")
            return True
        except Exception as e:
            logger.error(f"更新脚本 {script_name} 失败: {str(e)}")
            return False
    
    def delete_script(self, script_name):
        """
        删除脚本
        """
        if script_name not in self.scripts:
            logger.error(f"脚本 {script_name} 不存在")
            return False
        
        try:
            script_path = self.scripts[script_name]['path']
            os.remove(script_path)
            
            # 从内存中移除
            if script_name in self.scripts:
                del self.scripts[script_name]
            
            # 更新元数据
            self._save_metadata()
            logger.info(f"脚本 {script_name} 删除成功")
            return True
        except Exception as e:
            logger.error(f"删除脚本 {script_name} 失败: {str(e)}")
            return False
    
    def toggle_script_status(self, script_name, enabled):
        """
        启用或禁用脚本
        """
        if script_name not in self.scripts:
            logger.error(f"脚本 {script_name} 不存在")
            return False
        
        try:
            script = self.scripts[script_name]
            script['enabled'] = enabled
            
            # 更新脚本文件
            script_path = script['path']
            with open(script_path, 'r', encoding='utf-8') as f:
                content = f.read()
            
            # 替换ENABLED变量
            if 'ENABLED =' in content:
                content = re.sub(r'ENABLED\s*=\s*(True|False)', f'ENABLED = {enabled}', content)
            else:
                content += f'\n\n# 脚本状态控制\nENABLED = {enabled}\n'
            
            with open(script_path, 'w', encoding='utf-8') as f:
                f.write(content)
            
            # 更新元数据
            self._save_metadata()
            logger.info(f"脚本 {script_name} {'启用' if enabled else '禁用'}成功")
            return True
        except Exception as e:
            logger.error(f"切换脚本 {script_name} 状态失败: {str(e)}")
            return False
    
    def _save_metadata(self):
        """
        保存脚本元数据
        """
        try:
            metadata = {}
            for script_name, script_info in self.scripts.items():
                # 只保存可序列化的信息
                metadata[script_name] = {
                    'name': script_info['name'],
                    'description': script_info['description'],
                    'author': script_info['author'],
                    'version': script_info['version'],
                    'vulnerability_type': script_info['vulnerability_type'],
                    'severity': script_info['severity'],
                    'enabled': script_info['enabled'],
                    'required_params': script_info['required_params'],
                    'optional_params': script_info['optional_params']
                }
            
            with open(self.script_metadata_file, 'w', encoding='utf-8') as f:
                json.dump(metadata, f, indent=4, ensure_ascii=False)
        except Exception as e:
            logger.error(f"保存脚本元数据失败: {str(e)}")
    
    def validate_script(self, content):
        """
        验证脚本内容是否合法
        """
        try:
            # 创建临时模块进行验证
            import types
            temp_module = types.ModuleType('temp_script')
            exec(content, temp_module.__dict__)
            
            # 检查必要元素
            if not hasattr(temp_module, 'scan') or not callable(temp_module.scan):
                return False, "脚本必须包含scan函数"
            
            if not hasattr(temp_module, 'NAME'):
                return False, "脚本必须定义NAME变量"
            
            if not hasattr(temp_module, 'VULNERABILITY_TYPE'):
                return False, "脚本必须定义VULNERABILITY_TYPE变量"
            
            return True, "脚本验证通过"
        except Exception as e:
            return False, f"脚本验证失败: {str(e)}"
    
    def get_script_template(self, script_name, vuln_type):
        """
        获取脚本模板
        """
        # 使用字符串列表和join方法构建模板，避免嵌套三引号和复杂转义问题
        template_parts = [
            '# 漏洞扫描脚本模板\n\n',
            '# 脚本基本信息\n',
            f'NAME = "{script_name}"\n',
            f'DESCRIPTION = "{vuln_type}漏洞扫描脚本"\n',
            'AUTHOR = ""\n',
            'VERSION = "1.0"\n',
            'VULNERABILITY_TYPE = "placeholder_vulnerability_type"\n',
            'SEVERITY = "medium"  # low, medium, high, critical\n',
            'ENABLED = True\n\n',
            '# 参数定义\n',
            'REQUIRED_PARAMS = []\n',
            'OPTIONAL_PARAMS = {}\n\n',
            '# 导入必要的模块\n',
            'import requests\n',
            'import re\n',
            'from app.utils.logger import logger\n\n',
            'def scan(target, **kwargs):\n',
            '    """\n',
            '    执行漏洞扫描\n',
            '    \n',
            '    Args:\n',
            '        target: 扫描目标URL\n',
            '        **kwargs: 可选参数\n',
            '    \n',
            '    Returns:\n',
            '        dict: 扫描结果，格式如下:\n',
            '        {\n',
            '            \'vulnerable\': True/False,\n',
            '            \'details\': {},  # 漏洞详细信息\n',
            '            \'evidence\': \'\',  # 漏洞证据\n',
            '            \'remediation\': \'\'  # 修复建议\n',
            '        }\n',
            '    """\n',
            '    result = {\n',
            '        \'vulnerable\': False,\n',
            '        \'details\': {},\n',
            '        \'evidence\': \'\',\n',
            '        \'remediation\': \'\'\n',
            '    }\n',
            '    \n',
            '    try:\n',
            '        logger.info(f"开始扫描目标: {{target}}，漏洞类型: {{VULNERABILITY_TYPE}}")\n',
            '        \n',
            '        # TODO: 实现漏洞扫描逻辑\n',
            '        # 1. 发送请求\n',
            '        # 2. 分析响应\n',
            '        # 3. 判断是否存在漏洞\n',
            '        # 4. 收集证据\n',
            '        # 5. 提供修复建议\n',
            '        \n',
            '        # 示例代码\n',
            '        # response = requests.get(target, timeout=10)\n',
            '        # if vulnerability_found:\n',
            '        #     result[\'vulnerable\'] = True\n',
            '        #     result[\'details\'] = {\'key\': \'value\'}\n',
            '        #     result[\'evidence\'] = \'漏洞证据\'\n',
            '        #     result[\'remediation\'] = \'修复建议\'\n',
            '        \n',
            '    except Exception as e:\n',
            '        logger.error(f"扫描过程中发生错误: {{str(e)}}")\n',
            '    \n',
            '    return result\n'
        ]
        
        template = ''.join(template_parts)
        return template

# 创建全局脚本管理器实例
script_manager = VulnerabilityScriptManager()